from langchain.tools import tool, ToolRuntime
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

from model.Ark import ArkModel

def short_memory():

    @tool
    def get_weather(add: str, runtime: ToolRuntime):
        """ 获取一个地址的天气信息
            :Args:
                add: 地理位置
        """
        print('---> add is:', add, runtime.config)
        return '晴天'


    model = ArkModel().model
    agent = create_agent(
        model=model,
        tools=[get_weather],
        checkpointer=InMemorySaver() ## 添加检查点、可下载：langgraph-checkpoint-postgres、langgraph-checkpoint-redis 实现数据库存储
    )

    while True:
        try:
            user_input = input("请输入你的问题：>\n")
            response = agent.invoke({ 'messages': [{ 'role': 'user', 'content': user_input }] }, { 'configurable': { 'thread_id': '123123' } })
        except KeyboardInterrupt as e:
            break

def run():
    short_memory()